from models import  yolo
import torch
from PIL import Image,ImageDraw
from torchvision import transforms
# from utils import utils

# model = yolo.Model("models/yolov5s.yaml")
# model.load_state_dict(torch.load("models/yolov5s.pt"))
model = yolo.Model("yolov5s.yaml")
model.load_state_dict(torch.load("yolov5s.pt"))

model.eval()

img = Image.open(r"inference/images/zidane.jpg")
tf = transforms.Compose([
    transforms.Resize((512,640)),
    transforms.ToTensor()
])
img_tensor = tf(img)
pred = model(img_tensor[None])[0]
model.model[-1].export = True
torch.onnx.export(model, img_tensor[None], "my_yolov5s.onnx", verbose=True, opset_version=11, input_names=['images'],
                      output_names=['output1','output2','output3'])  # output_names=['classes', 'boxes']

